package edu.cn.nlsde.tmfst.EVA;

//import org.deeplearning4j.models.embeddings.wordvectors.WordVectors;

import java.io.BufferedReader;
import java.io.FileReader;
import java.io.IOException;
import java.util.*;
import java.util.Map.Entry;

public class CoherenceScore {

//    private WordVectors w2vModel = null;

    private HashMap<String, HashSet<Integer>> word2docset;

    public CoherenceScore(List<List<String>> docs) {
        word2docset = new HashMap<>();
        for (int docID = 0; docID < docs.size(); docID++) {
            List<String> terms = docs.get(docID);
            for (String term : terms) {
                HashSet<Integer> docSet;
                if (!word2docset.containsKey(term)) {
                    docSet = new HashSet<>();
                } else {
                    docSet = word2docset.get(term);
                }
                docSet.add(docID);
                word2docset.put(term, docSet);
            }
        }

    }

    public CoherenceScore(String path) throws IOException {
        word2docset = new HashMap<>();

        BufferedReader reader = new BufferedReader(new FileReader(path));

        String line;
        int docID = 0;
        try {
            line = reader.readLine();
            while (line != null) {
                String[] terms = line.trim().split("\\s+");
                for (String term : terms) {
                    HashSet<Integer> docSet;
                    if (!word2docset.containsKey(term)) {
                        docSet = new HashSet<Integer>();
                    } else {
                        docSet = word2docset.get(term);
                    }
                    docSet.add(docID);
                    word2docset.put(term, docSet);
                }
                line = reader.readLine();
                docID++;
            }
        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    public static void Caculate_Coherence(String trainPath, String wordmapPath, String phiPath, int topNum) throws IOException {
        CoherenceScore score = new CoherenceScore(trainPath);
        score.CoherenceScores(wordmapPath, phiPath, topNum);
    }

    public static void main(String args[]) throws IOException {
        String root = "G:/intellij/TopicModelForShortText/My_LDA/data4/new-tweet/";
        String kroot = root + "BTM_filter_0.1_0.01_20/1/";
        Caculate_Coherence(root + "new-tweet.data", kroot + "model-final.wordmap", kroot + "model-final.phi", 5);

//        String dataname = "new-tweet";
//        String method = "GSLDA";
//        String root = "G:/intellij/TopicModelForShortText/My_LDA/data/" + dataname + "_1000/wntm/";
//
//        String kroot = root +method+ "_0.1_0.01_50/0/";
//        Caculate_Coherence(root + dataname+".data",
//                kroot +"model-final.wordmap",
//                kroot +"model-final.phi",
//                5);


//        String[] nams = {
////                "news",
////                "new-tweet",
////                "dblp",
////                "yahooQA",
//                "20newsgroup",
////                "stackoverflow"
//        };
//        String[] methods = {
//                "LDA_0.1_0.01_",
////                "BTM_0.1_0.01_",
////                "unigram_",
////                "GSLDA",
//                "GSLDA_without_background_0.1_0.01_0.3_",
////                "GSLDA_without_background_0.1_0.01_5.0_100",
////                "GSLDA_without_background_0.1_0.01_5.0_100"
//        };
//        for (int topic_num = 60; topic_num < 110; topic_num += 20) {
//            for (String dataname : nams) {
//                String root = "G:/intellij/TopicModelForShortText/My_LDA/data3/" + dataname + "/wntm/";
//                for (String method : methods) {
////                    System.out.println(dataname + " " + method);
//                    System.out.println(dataname + " " + method+topic_num);
//                    for (int i = 0; i < 5; i++) {
//                        String kroot = root + method +topic_num+ "/" + i + "/";
//                        Caculate_Coherence(root + dataname + ".data",
//                                kroot + "model-final.wordmap",
//                                kroot + "model-final.phi",
//                                5);
//                    }
//                }
//            }
//        }
    }

    public double log(double value, double base) {
        return Math.log(value) / Math.log(base);
    }

    public double[][] loadPhi(String path) throws IOException {
        List<String> rows = new ArrayList<String>();

        BufferedReader reader = new BufferedReader(new FileReader(path));

        String line;
        try {
            line = reader.readLine();
            while (line != null) {
                rows.add(line.trim());
                line = reader.readLine();
            }
            reader.close();
        } catch (IOException e) {
            e.printStackTrace();
        }

        int l = rows.get(0).split(" ").length;
        int r = rows.size();

        double[][] phi = new double[r][l];

        for (int i = 0; i != r; i++) {
            String[] vals = rows.get(i).split(" ");
            for (int j = 0; j != l; j++) {
                phi[i][j] = Double.valueOf(vals[j]);
            }
        }

        return phi;
    }

    public HashMap<String, Integer> loadWordMap(String path) throws IOException {
        HashMap<String, Integer> wordmap = new HashMap<String, Integer>();

        BufferedReader reader = new BufferedReader(new FileReader(path));

        String line;
        try {
            line = reader.readLine();
            while (line != null) {
                if (line.split("\\s+")[0].equals("#virtualWord#")) {
                    line = reader.readLine();
                    continue;
                }
                wordmap.put(line.split("\\s+")[0], Integer.valueOf(line.split("\\s+")[1]));
                line = reader.readLine();
            }
            reader.close();
        } catch (IOException e) {
            e.printStackTrace();
        }
        return wordmap;
    }

    public ArrayList<List<Entry<String, Double>>> sorted_topicwords(HashMap<String, Integer> wordmap,
                                                                    double[][] phi) {
        ArrayList<List<Entry<String, Double>>> res = new ArrayList<List<Entry<String, Double>>>();

        int K = phi.length;

        for (int k = 0; k != K; k++) {
            HashMap<String, Double> term2weight = new HashMap<String, Double>();
            for (String term : wordmap.keySet())
                term2weight.put(term, phi[k][wordmap.get(term)]);

            List<Entry<String, Double>> pairs = new ArrayList<Entry<String, Double>>(
                    term2weight.entrySet());
            Collections.sort(pairs, new Comparator<Entry<String, Double>>() {
                public int compare(Entry<String, Double> o1,
                                   Entry<String, Double> o2) {
                    return (o2.getValue().compareTo(o1.getValue()));
                }
            });

            res.add(pairs);
        }

        return res;
    }

    public void print_topics(String wordmapPath, String phiPath, int top_n) throws IOException {
        HashMap<String, Integer> wordmap = this.loadWordMap(wordmapPath);
        double[][] phi = this.loadPhi(phiPath);
        ArrayList<List<Entry<String, Double>>> pairs_list = this
                .sorted_topicwords(wordmap, phi);

        int K = phi.length;
        for (int k = 0; k != K; k++) {
            System.out.println("Topic " + k + ":");
            for (int i = 0; i != top_n; i++) {
                System.out.println(pairs_list.get(k).get(i).getKey() + " "
                        + pairs_list.get(k).get(i).getValue());
            }
        }
    }

//    public void CoherenceScores(String topicWordPath, int top_n) throws IOException {
//
//        ArrayList<ArrayList<String>> sorted_pairs = LDAFileProcess.read_topic_words(topicWordPath, 20);
//        int K = sorted_pairs.size();
//        double t_score = 0.0;
//        for (int k = 0; k != K; k++) {
//            double score = 0.0;
//            for (int i = 1; i != top_n; i++) {
//                for (int j = 0; j != i; j++) {
//                    //System.out.println(sorted_pairs.get(0).get(i)+" "+sorted_pairs.get(k).get(j));
//                    score += this.getCoherenceScore(sorted_pairs.get(k).get(i),
//                            sorted_pairs.get(k).get(j));
//                }
//            }
//            //score = score * 2.0 / top_n * (top_n - 1.0);
//            t_score += score;
//            //System.out.println("topic " + k + ": " +score);
//        }
//        System.out.println(t_score / K);
//    }

    public double getCoherenceScore(String w1, String w2) {
        HashSet<Integer> set1 = word2docset.get(w1);
        HashSet<Integer> set2 = word2docset.get(w2);
        //System.out.println(w2);
        double size2 = set2.size();
        double jointSize = 0.0;
        for (Integer id : set1) {
            if (set2.contains(id))
                jointSize++;
        }
        return log((jointSize + (1E-12)) / size2, 2);
    }

//    public static void Caculate_Coherence(String trainPath, String topicWordPath, int topNum) throws IOException {
//        CoherenceScore score = new CoherenceScore(trainPath);
//        score.CoherenceScores(topicWordPath, topNum);
//    }

    public void CoherenceScores(String wordmapPath, String phiPath, int top_n) throws IOException {
        HashMap<String, Integer> wordmap = this.loadWordMap(wordmapPath);
        double[][] phi = this.loadPhi(phiPath);

        ArrayList<List<Entry<String, Double>>> sorted_pairs = this.sorted_topicwords(wordmap, phi);
        int K = sorted_pairs.size();
        double t_score = 0.0;
        for (int k = 0; k != K; k++) {
            double score = 0.0;
            for (int i = 1; i != top_n; i++) {
                for (int j = 0; j != i; j++) {
                    score += this.getCoherenceScore(sorted_pairs.get(k).get(i).getKey(),
                            sorted_pairs.get(k).get(j).getKey());
                }
            }
            t_score += score;
//			System.out.println("topic " + k + ": " +score);
            //System.out.println(score);
        }
        System.out.println(t_score / K);
    }


    public double[] CoherenceScores(HashMap<String, Integer> wordmap, double[][] phi, int top_n) {
        ArrayList<List<Entry<String, Double>>> sorted_pairs = this.sorted_topicwords(wordmap, phi);
        int K = sorted_pairs.size();
        double[] result = new double[K];
        double t_score = 0.0;
        for (int k = 0; k != K; k++) {
            double score = 0.0;
            for (int i = 1; i != top_n; i++) {
                for (int j = 0; j != i; j++) {
                    score += this.getCoherenceScore(sorted_pairs.get(k).get(i).getKey(),
                            sorted_pairs.get(k).get(j).getKey());
                }
            }
            result[k] = score;
            t_score += score;
//			System.out.println("topic " + k + ": " +score);
            //System.out.println(score);
        }
        System.out.println(t_score / K);
        return result;
    }
}
